
/**
 * This is tool for testing tcp server performance 
 */

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <assert.h>
#include <time.h>
#include <string.h>

#ifdef _WIN32
# include <winsock2.h>
# define strcasecmp(s, t) _stricmp(s, t)
# define strncasecmp(s, t,n) _strnicmp(s, t, n)
# define sleep(n) Sleep((n)*1000)
# define usleep(n) Sleep((n))
#else
# include <unistd.h>
# include <netinet/in.h>
# include <pthread.h>
#endif

#include "aio_base.h"
#include "aio_log.h"
#include "aio_socket.h"
#include "aio_slot.h"
#include "getopt.h"

/* globle variable define */
static char server_address[64] = {0};
static char dstip[16] = {0};
static unsigned short dstport = 0;
static int thread_counts = 1;
static int connection_counts_per_thread = 1;
static int connect_timeout = 5;
static int echo_packet_size = 100;
static int echo_frequency = 1000;
static char loadfiledir[128] = {0};
static int receive_print = 0;
static int log_level = 2;
static char *sendbuf = NULL;
static int exit_time = 100;
const char* helpinfo = 
       "\nThis is tool for testing tcp server performance.\n\n"
       "Usage: benchmark [OPTIONS]\n"
       "    -h, display this message then exit.\n"
       "    -v, display version information then exit.\n"
       "    -d, debug log onoff\n"
       "    -p, print receive packet onoff\n"
       "    -s, server address, [ip:port]\n"
       "    -n, thread counts, default value:1\n"
       "    -c, tcp connection counts per thread, default value:1\n"
       "    -o, tcp connect timeout time [second], default value:10s\n"
       "    -l, echo tcp packet size [bytes], default value:100\n" 
       "    -t, echo frequency per tcp connection [millisecond], default value:1000\n"
       "    -f, open file name that you want to send\n"
       "    -e, exit program after [n] seconds. default: 100s\n"
       "Report bugs to <hlxxxx@google.com>\n";

/* data collect */
struct thread_data {
        unsigned conn_count;
        unsigned lost_conn_count;
        unsigned send_bytes; 
        unsigned recv_bytes;
};

void
onclosed(struct aio_slot *slot)
{
        struct thread_data *td;
        struct aio_timer *timer;
        debugp("Enter onclosed function... \n");
        assert(slot);
     
        td = (struct thread_data *)ABUSRP(SLOTABP(slot)); 
        assert(td);

        td->lost_conn_count++;

        timer = (struct aio_timer *)SLOTUSRP(slot);

        if (timer) {
           aio_timer_cancel(timer);
        }    
}

void
ontimer(struct aio_timer *at, void *arg)
{
        struct aio_buf *buf;
        struct aio_slot *slot;
        debugp("Enter ontimer1 function.\n");

        slot = (struct aio_slot*)arg;
        assert(slot);

        if (echo_packet_size > 0) {
                struct thread_data *td;
                buf = AIO_BUF_ALLOC(SLOTABP(slot));
                assert(buf);

                if (sendbuf) {
                        memcpy(BPTR(buf), sendbuf, echo_packet_size);
                }

                td = (struct thread_data*)ABUSRP(SLOTABP(slot)); 
                assert(td);
                td->send_bytes += echo_packet_size; 
      
                BLEN(buf) = echo_packet_size;
                AIO_WRITE(slot, buf, NULL);
        }
}

void    
onreceived(struct aio_slot *slot) 
{
        struct aio_buf *rbuf;
        struct thread_data *td;

        debugp("Enter onread function...\n");
        assert(slot);

        if (SLOTIOERR(slot) != 0) {
                debugp("Slot operate timeout or other error.\n");
                aio_slot_close(slot, 1);
                return;
        }

        rbuf = AIO_READ(slot, 1);
        assert(rbuf);

        td = (struct thread_data *)ABUSRP(SLOTABP(slot));
        assert(td);
        td->recv_bytes += BLEN(rbuf);

        if (receive_print) {
                *BEND(rbuf) = 0;
                printf("%s\n", BPTR(rbuf));
        }

        if (echo_frequency == 0) {
                td->send_bytes += echo_packet_size; 
                AIO_WRITE(slot, rbuf, NULL);
        } else {
                AIO_BUF_FREE(rbuf);
        }
}

void
onconnected(struct aio_slot *slot)
{
        struct thread_data *td;
        struct aio_timer *timer;
        debugp("Enter onconnected function...\n");

        assert(slot);

        if (SLOTIOERR(slot) != 0) {
                log_warn("Slot operate timeout or other error.\n");
                aio_slot_close(slot, 1);
                return;
        }

        td = (struct thread_data *)ABUSRP(SLOTABP(slot));
        td->conn_count++;
        
        if (aio_setread(slot, onreceived, 1, 0) < 0) {
                log_warn("Failed to set aio slot read.\n");
                aio_slot_close(slot, 1);
        }

        if (echo_frequency == 0) {
                ontimer(NULL, slot); return;
        }
  
        timer = aio_timer_set(SLOTABP(slot), echo_frequency, ontimer, slot);
        assert(timer);
        SLOTUSRP(slot) = timer;
}

#ifdef _WIN32
DWORD WINAPI
thread_proc(LPVOID arg)
#else
void*
thread_proc(void *arg)
#endif
{
        int i;
        struct aio_base *base;
        struct aio_slot *slot;

        base = aio_create(NULL, connection_counts_per_thread+10);
        if (base == NULL) {
                log_warn("Failed to create AIO.\n");
                goto exit;
        }
 
        ABUSRP(base) = arg;

        for (i = 0; i < connection_counts_per_thread; i++)
	{
		slot = aio_slot_open(base, NULL, 0, AIO_TCP);
		if (slot == NULL) {
			log_warn("Failed to open a new slot.\n");
                        goto exit;
		}

                AIO_SETDSTADDR2(slot, inet_addr(dstip), dstport);
		aio_setclose(slot, onclosed);
        
                if (aio_connect(slot, onconnected, connect_timeout*1000) < 0) {
                        log_warn("Failed to call aio_connect function.\n");
                        goto exit;
                }
        }

        aio_loop(base, exit_time*1000); 
        debugp("The thread exit.\n"); 
exit:
#ifdef _WIN32
        return 0;
#else
        return;
#endif
}

static int 
parse_options(int argc, char *argv[])
{
	int c;
	extern char *optarg;
        extern int optind, opterr;

        opterr = 0;

        if (argc == 1) return -1;

	while ((c = getopt(argc, argv, "hvpds:n:c:l:o:t:f:e:")) != -1) {
		switch(c) {
                case 'h':
                        printf(helpinfo);
                        return 1;
                case 'v':
                        printf("benchmark-0.0.1\n");
                        return 1;
                case 'p':
                        receive_print = 1;
                        break;
                case 'd':
                        log_level = 0;
                        break;
                case 's':
                        if (strlen(optarg) > 64) return -1;
                        strcpy(server_address, optarg);
                        break;
                case 'n':
                        thread_counts = atoi(optarg);
                        break;
                case 'c':
                        connection_counts_per_thread = atoi(optarg);
                        break;
                case 'o':
                        connect_timeout = atoi(optarg);
                        break;
                case 'l':
                        echo_packet_size = atoi(optarg);
                        break;
                case 't':
                        echo_frequency = atoi(optarg);
                        break;
                case 'f':
                        if (strlen(optarg) > 128) return -1;
                        strcpy(loadfiledir, optarg);
                        break;
                case 'e':
                        exit_time = atoi(optarg); 
                        break;
                default:
                        printf("Unknown input arguments [%c].\n", c);
                        return -1;
                }
        }

        if (optind < argc) {
                printf("Unexpected argument %s\n", argv[optind]);
                return -1;
        }

	return 0;
}

int main(int argc, char** argv)
{
        int i, ret;
        char* pos;
        struct thread_data **pptd; 
        unsigned total_conn_count = 0; 
        unsigned total_lost_conn_count = 0;
        unsigned total_send_bytes = 0;
        unsigned total_recv_bytes = 0;

#ifdef _WIN32
        WSADATA wsaData;        
        WORD version = MAKEWORD(2, 0);        
        WSAStartup(version, &wsaData);
#else
        pthread_t tid;
#endif

        ret = parse_options(argc, argv);
        if (ret) {
                if (ret == -1) printf(helpinfo);
                exit(-1);
        }
               
        printf("\n------Benchmark start infomation------\n"
               "server address: %s\n"
               "thread counts: %d\n"
               "connection counts per thread: %d\n" 
               "connect timeout time: %ds\n"
               "echo packet size: %d bytes\n"
               "echo frequency: %d ms\n"
               "the file dir: %s\n"
               "-------Exit Ctrl+C-------------------\n\n",
               server_address, 
               thread_counts, 
               connection_counts_per_thread,
               connect_timeout,
               echo_packet_size,
               echo_frequency,
               loadfiledir);

        pos = strstr(server_address, ":"); 
        if (pos == NULL) {
               log_err("Invalid server address infomation.\n");
               exit(-1);
        }

        strncpy(dstip, server_address, pos - server_address);
        dstport = atoi(++pos);

        if (strlen(dstip) <= 0 || dstport <= 0) {
                log_err("Invalid server ip or port.\n");
                exit(-1);
        }

        if (strlen(loadfiledir) > 0) {
                FILE* fp = fopen(loadfiledir, "r");
                if (fp == NULL) {
                        log_err("Failed to open file %s\n", loadfiledir);
                        exit(-1);
                }
                sendbuf = calloc(1, AIO_BUF_PLSIZE);
                echo_packet_size = (int)fread(sendbuf, 1, AIO_BUF_PLSIZE, fp);

                if (echo_packet_size <= 0) {
                        log_err("Failed to read %s file\n", loadfiledir);
                        exit(-1);
                }

                printf("File Content: \n\n%s\n\n", sendbuf);
                fclose(fp);
        }
                        
        log_set_silence_level(log_level);

        pptd = (struct thread_data**)calloc(thread_counts, sizeof(void*));

        for (i = 0; i < thread_counts; i++) {
                pptd[i] = (struct thread_data*)calloc(1, sizeof(struct thread_data));
                assert(pptd[i]);
#ifdef _WIN32
                if (!CreateThread(NULL, 0, thread_proc, (LPVOID)pptd[i], 0, 0)) {
#else
                if(pthread_create(&tid, NULL, thread_proc, pptd[i])) {
#endif
                        log_err("Failed to create thread.\n");
                        exit(-1);
                }
        }

        { /* waiting for testing */
                int ntimes = exit_time;
                printf("Testing"); 
                while (ntimes--) {
                        sleep(1); 
                        printf(".");
                        fflush(stdout);
                }
                printf("\n");
        }

        { /* waiting for worker thread exit */
                int ntimes = 6;
                printf("Data collecting"); 
                while (ntimes--) {
                        sleep(1); 
                        printf(".");
                        fflush(stdout);
                }
                printf("\n");
        }
                     
        for (i = 0; i < thread_counts; i++) {
                total_conn_count += pptd[i]->conn_count; 
                total_lost_conn_count += pptd[i]->lost_conn_count;
                total_send_bytes += pptd[i]->send_bytes;
                total_recv_bytes += pptd[i]->recv_bytes;
                free(pptd[i]);
        }

        printf("total time: %ds\n"
               "total thread count: %d\n"
               "total connection count: %d\n"
               "success connection count: %d\n"
               "lost connection count: %d\n"
               "total send bytes: %d\n"
               "total recv bytes: %d\n"
               "average bandwith: %d bytes/second\n",
               exit_time, thread_counts, 
               thread_counts*connection_counts_per_thread,
               total_conn_count, total_lost_conn_count, 
               total_send_bytes, total_recv_bytes,
               total_recv_bytes/exit_time);

        if (sendbuf) free(sendbuf);
        if (pptd) free(pptd);


        printf("Echo client exit.\n");
        exit(0);
}

